為了解說方便，讓範例程式的結構更完整，老頭寫了第二個範例程式（程式的 colab 檔由此下載），這個範例程式更完整的呈現了用 Flax 來設計、訓練、儲存及載入模型的一系列流程，包含了：

• 載入所需的程式庫：常用的函式庫，及其慣用的 import 語法。
• 載入資料集：使用 TensorFlow 的 TFDS 服務，載入 cifar10 資料集。
• 自訂模型：用 Flax 定義了兩種模型範例
  • MLP 模型
  • CNN 模型
• 輔助函式：
  • 損失函式
  • 參數最佳化方法
  • 輔助評估函式
• 訓練準備
  • 單一 batch 訓練函式
  • 單一 epoch 訓練函式
  • 模型評估函式
• 執行訓練：分別訓練前面定義的 MLP 及 CNN 模型。
• 儲存及載入模型：說明如何在模型參數存檔，及由模型參數重建模型。

除了方便解說之外，老頭也想留下一個樣板 (template) ，如果有機會用 Flax 解決類似的問題，就直接拿它來改，不必重頭再寫一次。

接下來，老頭將會一一介紹這個範例程式的各個部分。